package org.odk.collect.android.utilities;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.NoSuchAlgorithmException;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.spec.X509EncodedKeySpec;
import java.util.Random;
import java.util.UUID;
import java.util.Vector;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import org.javarosa.core.model.instance.TreeElement;
import org.kxml2.io.KXmlParser;
import org.odk.collect.android.application.Collect;
import org.odk.collect.android.logic.FormController;
import org.odk.collect.android.tasks.InstanceUploaderTask;
import org.xmlpull.v1.XmlPullParser;
import org.xmlpull.v1.XmlPullParserException;

import android.util.Base64;
import android.widget.Toast;

public class EncryptionUtils {

	private static final String DEFAULT_ALGORITHM = "AES/CFB/NoPadding";
	private static final String ENC_FILE_EXTENSION = ".enc";
	private static final int DEFAULT_KEY_LENGTH = 16; // 128 bits
	private static final String RSA = "RSA";
	private static final String UTF_8 = "UTF-8";
	private static final String SYMMETRIC_KEY = "symmetricKey";
	private static final String AES = "AES";
	private static final String BASE64_PIN_ENCRYPTED_SYMMETRIC_KEY = "base64pinEncryptedSymmetricKey";
	private static final String BASE64_RSA_ENCRYPTED_SYMMETRIC_KEY = "base64RsaEncryptedSymmetricKey";
	private static final String ALGORITHM = "algorithm";
	private static final String DATA = "data";
	private static final String ID = "id";
	private static final String VERSION = "version";
	private static final String UI_VERSION = "uiVersion";
	private static final String BASE64_RSA_PUBLIC_KEY = "base64RsaPublicKey";
	private static final String ENCRYPTED = "encrypted";
	private static final String ENCRYPTED_XML_FILE = "encryptedXmlFile";
	private static final String META = "meta";
	private static final String INSTANCE_ID = "instanceID";
	private static final String TIME_START = "timeStart";
	private static final String TIME_END = "timeEnd";
	private static final String MEDIA = "media";
	private static final String NAME = "name";
	private static final String FILE = "file";
	
	private static void emitAttribute(StringBuilder b, String name, String value) {
		if ( value == null ) return;
		b.append(name);
		b.append("=\"");
		b.append(value);
		b.append("\" ");
	}
	
	private static void emitTag(StringBuilder b, String name, String value) {
		if ( value != null ) {
			b.append("<" + name + ">");
			b.append(value);
			b.append("</" + name + ">");
		} else {
			b.append("<" + name + "/>");
		}
	}
	private static Cipher getPinCipher(long pin, boolean forEncryption) throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException {
    	Random rPin = new Random(pin);
    	byte[] pinKey = new byte[16];
    	rPin.nextBytes(pinKey);
    	Cipher c = Cipher.getInstance(AES);
    	SecretKeySpec kPin = new SecretKeySpec(pinKey,AES);
    	if ( forEncryption ) {
    		c.init(Cipher.ENCRYPT_MODE, kPin);
    	} else {
    		c.init(Cipher.DECRYPT_MODE, kPin);
    	}
    	return c;
	}
	
	private static byte[] generatePinEncryptedKey(long pin, int keyLength) throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException {
    	Cipher c = getPinCipher(pin, true);
    	// generate key
    	SecureRandom r = new SecureRandom();
    	byte[] key = new byte[keyLength];
    	r.nextBytes(key);
    	// encrypt this key...
    	byte[] pinEncryptedKey = c.doFinal(key);
    	return pinEncryptedKey;
	}
	
	private static byte[] decodePinEncryptedBase64Value(long pin, String value) throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException {
		Cipher c = getPinCipher(pin, false);
    	// decrypt the incoming Base64 value...
    	byte[] pinDecryptedKey = c.doFinal(Base64.decode(value, Base64.NO_WRAP));
    	return pinDecryptedKey;
	}
	
	private static Integer getVersion(FormController fc) {
		String value = fc.getRootElementAttribute(VERSION);
    	if ( value == null || value.length() == 0 ) {
    		return null;
    	}
    	return Integer.parseInt(value);
	}
	
	private static Integer getUiVersion(FormController fc) {
		String value = fc.getRootElementAttribute(UI_VERSION);
    	if ( value == null || value.length() == 0 ) {
    		return null;
    	}
    	return Integer.parseInt(value);
	}
	
	private static byte[] getPublicEncryptionKey(FormController fc) {
		String value = fc.getRootElementAttribute(BASE64_RSA_PUBLIC_KEY);
    	if ( value == null || value.length() == 0 ) {
    		return null;
    	}
    	return Base64.decode(value, Base64.DEFAULT);
	}
	
	public static Long askForPin() {
    	Toast t = Toast.makeText(Collect.getInstance().getApplicationContext(), "Ask for PIN", Toast.LENGTH_LONG);
    	t.show();
    	String PIN = "1234";
    	long pin = Long.parseLong(PIN);
    	return pin;
	}
	
	public static Long maybeAskForPin(String instancePath) {
		if ( instancePath == null ) return null;
		String infoFilePath = instancePath + ".info";
    	File f = new File(infoFilePath);

    	if (f.exists() ) {
    		return askForPin();
    	}
    	return null;
	}
	
	/**
	 * @param fc
	 * @param instancePath
	 * @return pin if we are encrpyting, null otherwise...
	 * @throws NoSuchAlgorithmException
	 * @throws NoSuchPaddingException
	 * @throws IllegalBlockSizeException
	 * @throws BadPaddingException
	 * @throws InvalidKeyException
	 * @throws IOException
	 */
	public static Long configureEncryption(FormController fc, String instancePath) throws NoSuchAlgorithmException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException, InvalidKeyException, IOException {
        
        byte[] publicKey = getPublicEncryptionKey(fc);
        if ( publicKey == null ) {
        	return null;
        }
        String algorithm = DEFAULT_ALGORITHM;
        long pin = askForPin();
    	byte[] pinEncryptedKey = generatePinEncryptedKey(pin, DEFAULT_KEY_LENGTH);
    	
    	StringBuilder b = new StringBuilder();
    	b.append("<" + SYMMETRIC_KEY + " ");
    	emitAttribute(b, BASE64_RSA_PUBLIC_KEY, 
    			Base64.encodeToString(publicKey, Base64.NO_WRAP));
    	emitAttribute(b, ALGORITHM, algorithm);
    	emitAttribute(b, BASE64_PIN_ENCRYPTED_SYMMETRIC_KEY, 
    			Base64.encodeToString(pinEncryptedKey, Base64.NO_WRAP));
    	b.append(" />");
    	
    	String infoFilePath = instancePath + ".info";
    	File f = new File(infoFilePath);
    	FileOutputStream out = new FileOutputStream(f);
    	OutputStreamWriter writer = new OutputStreamWriter(out, UTF_8);
    	writer.write(b.toString());
    	writer.close();
		return pin;
	}
	
	public static Cipher getFileCipher(Long pin, String instancePath, boolean forEncryption) throws IOException, XmlPullParserException, InvalidKeyException, NoSuchAlgorithmException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException, InvalidAlgorithmParameterException {
		if ( pin == null ) return null;
		
		String infoFilePath = instancePath + ".info";
		File f = new File(infoFilePath);
		if ( !f.exists() ) {
			throw new IllegalStateException("Expected encryption information file not found " + infoFilePath);
		}
		
		FileInputStream in = new FileInputStream(f);
		KXmlParser p = new KXmlParser();
		p.setInput(in, UTF_8);
		p.nextTag();
		p.require(XmlPullParser.START_TAG, null, SYMMETRIC_KEY);
		String algorithm = p.getAttributeValue(null, ALGORITHM);
		String pinEncryptedBase64Value = p.getAttributeValue(null, BASE64_PIN_ENCRYPTED_SYMMETRIC_KEY); 
		in.close();

		byte[] key = decodePinEncryptedBase64Value(pin, pinEncryptedBase64Value);
    	Cipher c = Cipher.getInstance(algorithm);
    	SecretKeySpec kPin = new SecretKeySpec(key,algorithm);
    	if ( forEncryption ) {
    		SecureRandom r = new SecureRandom();
    		byte[] iv = new byte[key.length];
    		r.nextBytes(iv);
        	IvParameterSpec baseIv = new IvParameterSpec(iv);
    		c.init(Cipher.ENCRYPT_MODE, kPin, baseIv);
    	} else {
    		byte[] iv = new byte[key.length];
        	IvParameterSpec baseIv = new IvParameterSpec(iv);
    		c.init(Cipher.DECRYPT_MODE, kPin, baseIv );
    	}
    	return c;
	}
	
	public static void emitSubmissionManifest(Long pin, FormController fc, String instancePath) {
		if ( pin == null ) return;

		String infoFilePath = instancePath + ".info";
    	String envelopeFilePath = instancePath + ".envelope";

    	File instanceFile = new File(instancePath);
		File submissionFile = new File(instanceFile.getParent(), InstanceUploaderTask.XML_SUBMISSION_FILENAME);
		File infoFile = new File(infoFilePath);
    	File envelopeFile = new File(envelopeFilePath);
    	
		try {
			if ( !infoFile.exists() ) {
				throw new IllegalStateException("Expected encryption information file not found " + infoFilePath);
			}

			String algorithm;
			String base64RsaPublicKey;
			String base64PinEncryptedKeyValue;
			FileInputStream in = null;
			try {
				in = new FileInputStream(infoFile);
				KXmlParser p = new KXmlParser();
				p.setInput(in, UTF_8);
				p.nextTag();
				p.require(XmlPullParser.START_TAG, null, SYMMETRIC_KEY);
				algorithm = p.getAttributeValue(null, ALGORITHM);
				base64RsaPublicKey = p.getAttributeValue(null, BASE64_RSA_PUBLIC_KEY);
				base64PinEncryptedKeyValue = p.getAttributeValue(null, BASE64_PIN_ENCRYPTED_SYMMETRIC_KEY); 
				in.close();
			} catch ( Exception e) {
				if ( in != null ) {
					in.close();
				}
				throw new IllegalStateException("Unable to read " + infoFile.getAbsolutePath());
			}

		byte[] key = decodePinEncryptedBase64Value(pin, base64PinEncryptedKeyValue);
		
	    // create public key
	    X509EncodedKeySpec publicKeySpec = new X509EncodedKeySpec(
	    		Base64.decode(base64RsaPublicKey, Base64.NO_WRAP));
	    KeyFactory kf = KeyFactory.getInstance(RSA);
	    PublicKey pk = kf.generatePublic(publicKeySpec);
	    
	    // create RSA public key cipher
	    Cipher pkCipher = Cipher.getInstance(RSA);
	    
	    // write AES key
	    pkCipher.init(Cipher.ENCRYPT_MODE, pk);
	    byte[] pkEncryptedKey = pkCipher.doFinal(key);

    	String formId = fc.getFormId();
    	Integer version = getVersion(fc);
    	Integer uiVersion = getUiVersion(fc);

    	//////////////////////////////////
	    // now figure out what goes in the envelope...
	    File dir = instanceFile.getParentFile();
    	
    	StringBuilder b = new StringBuilder();
    	b.append("<" + DATA + " ");
    	emitAttribute(b, ID, formId);
    	emitAttribute(b, VERSION, (version == null) ? null : Integer.toString(version));
    	emitAttribute(b, UI_VERSION, (uiVersion == null) ? null : Integer.toString(uiVersion));
    	emitAttribute(b, ENCRYPTED, "yes");
    	b.append(">");
    	b.append("<" + META + ">");
    	TreeElement e = fc.getSubmissionMetaBlock();
    	if ( e != null ) {
    		Vector<TreeElement> v;
    		// timeStart...
    		v = e.getChildrenWithName(TIME_START);
    		if ( v.size() == 1 ) {
    			emitTag(b, TIME_START, v.get(0).getValue().toString());
    		}
    		// timeEnd...
    		v = e.getChildrenWithName(TIME_END);
    		if ( v.size() == 1 ) {
    			emitTag(b, TIME_END, v.get(0).getValue().toString());
    		}
    		// instance id...
    		v = e.getChildrenWithName(INSTANCE_ID);
    		if ( v.size() == 1 ) {
    			emitTag(b, INSTANCE_ID, v.get(0).getValue().toString());
    		} else {
    			// missing or two or more!!!
    			emitTag(b, INSTANCE_ID, "uuid:" + UUID.randomUUID().toString());
    		}
    	} else {
			// missing meta block!!!
			emitTag(b, INSTANCE_ID, "uuid:" + UUID.randomUUID().toString());
    	}
    	b.append("</" + META + ">");
    	emitTag(b, ALGORITHM, algorithm);
    	emitTag(b, BASE64_RSA_PUBLIC_KEY, base64RsaPublicKey);
    	emitTag(b, BASE64_RSA_ENCRYPTED_SYMMETRIC_KEY, 
    			Base64.encodeToString(pkEncryptedKey, Base64.NO_WRAP));
    	emitTag(b, ENCRYPTED_XML_FILE, submissionFile.getName());
    	b.append("\n");
    	File[] allFiles = dir.listFiles();
    	for ( File file : allFiles ) {
    		String name = file.getName();
    		if ( name.equals(instanceFile.getName()) ) continue;
    		if ( name.equals(submissionFile.getName()) ) continue;
    		if ( name.equals(infoFile.getName()) ) continue;
    		if ( name.equals(envelopeFile.getName()) ) continue;
    		b.append("<" + MEDIA + ">");
    		if ( name.endsWith(ENC_FILE_EXTENSION) ) {
            	emitTag(b, NAME, name.substring(0,name.length()-ENC_FILE_EXTENSION.length()));
            	emitTag(b, ENCRYPTED, "yes");
            	emitTag(b, FILE, file.getName());
    		} else {
	        	emitTag(b, NAME, file.getName());
	        	emitTag(b, ENCRYPTED, "no");
	        	emitTag(b, FILE, file.getName());
    		}
    		b.append("</" + MEDIA + ">\n");
    	}
    	b.append("</" + DATA + ">\n");
    	
    	FileOutputStream out = new FileOutputStream(envelopeFile);
    	OutputStreamWriter writer = new OutputStreamWriter(out, UTF_8);

    	writer.write(b.toString());
    	writer.close();
		} catch (Exception e) {
			e.printStackTrace();
			throw new IllegalStateException("unexpected error producing envelope file");
		}
	}
}
